package org.example.myframework.frame.common.websocket.config;

import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import javax.websocket.HandshakeResponse;
import javax.websocket.server.HandshakeRequest;
import javax.websocket.server.ServerEndpointConfig;
import java.util.List;
import java.util.Map;

/**
 * 鉴权
 */
@Component
public class WebSocketServerConfig extends ServerEndpointConfig.Configurator {

    @Override
    public boolean checkOrigin(String originHeaderValue) {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        return true;
    }

    @Override
    public void modifyHandshake(ServerEndpointConfig sec, HandshakeRequest request,
                                HandshakeResponse response) {
        // 获取所有ws请求连接的参数
        Map<String, List<String>> parameterMap = request.getParameterMap();
        List<String> erp = parameterMap.get("erp");
        if (!CollectionUtils.isEmpty(erp)) {
            sec.getUserProperties().put("erp", erp.get(0));
        }
    }
}
